#include "stdafx.h"
#include "net/SMTPClientSession.hpp"
#include "net/MailStream.hpp"
#include "net/SocketStream.hpp"
#include "filesys/environment.hpp"
using JHCPP::filesys::CEnvironment;
#include "stream/base64decode.hpp"
#include "stream/base64encode.hpp"
#include "stream/StreamCopier.hpp"
using JHCPP::stream::CBase64Encode;
using JHCPP::stream::CBase64Decode;
using JHCPP::stream::CStreamCopier;
#include "code/hmacengine.hpp"
#include "code/md5engine.hpp"
#include "code/sha1engine.hpp"
#include <sstream>
#include <fstream>
#include <iostream>

NET_NAMESPACE_BEGIN

CSMTPClientSession::CSMTPClientSession(const CStreamSocket& socket) : m_socket(socket), m_isOpen(false)
{
}

CSMTPClientSession::CSMTPClientSession(const std::string& host, UInt16 port)
	: m_socket(CSocketAddress(host, port)), m_isOpen(false)
{
}

CSMTPClientSession::~CSMTPClientSession()
{
	try
	{
		close();
	}
	catch (...)
	{
	}
}

void CSMTPClientSession::setTimeout(const CTimeSpan& timeout)
{
	m_socket.setReceiveTimeout(timeout);
}

CTimeSpan CSMTPClientSession::getTimeout() const
{
	return m_socket.getReceiveTimeout();
}

void CSMTPClientSession::login(const std::string& hostname, std::string& response)
{
	open();
	int status = sendCommand("EHLO", hostname, response);
	if (isPermanentNegative(status))
		status = sendCommand("HELO", hostname, response);
	if (!isPositiveCompletion(status)) throw SMTPException("Login failed", response, status);
}

void CSMTPClientSession::login(const std::string& hostname)
{
	std::string response;
	login(hostname, response);
}

void CSMTPClientSession::login()
{
	login(CEnvironment::nodeName());
}

void CSMTPClientSession::loginUsingCRAMMD5(const std::string& username, const std::string& password)
{
	THMACEngine<CMD5Engine> hmac(password);
	loginUsingCRAM(username, "CRAM-MD5", hmac);
}

void CSMTPClientSession::loginUsingCRAMSHA1(const std::string& username, const std::string& password)
{
	THMACEngine<CSHA1Engine> hmac(password);
	loginUsingCRAM(username, "CRAM-SHA1", hmac);
}

void CSMTPClientSession::loginUsingCRAM(const std::string& username, const std::string& method, CDigestEngine& hmac)
{
	std::string response;
	int status = sendCommand(std::string("AUTH ") + method, response);

	if (!isPositiveIntermediate(status)) throw SMTPException(std::string("Cannot authenticate using ") + method, response, status);
	std::string challengeBase64 = response.substr(4);

	std::istringstream istr(challengeBase64);
	CBase64Decode decoder(istr);
	std::string challenge;
	CStreamCopier::copyToString(decoder, challenge);

	hmac.update(challenge);

	const CDigestEngine::Digest& digest = hmac.digest();
	std::string digestString(CDigestEngine::digestToHex(digest));

	std::string challengeResponse = username + " " + digestString;

	std::ostringstream challengeResponseBase64;
	CBase64Encode encoder(challengeResponseBase64);
	encoder << challengeResponse;
	encoder.close();

	status = sendCommand(challengeResponseBase64.str(), response);
	if (!isPositiveCompletion(status)) throw SMTPException(std::string("Login using ") + method + " failed", response, status);  
}

void CSMTPClientSession::loginUsingLogin(const std::string& username, const std::string& password)
{
	std::string response;
	int status = sendCommand("AUTH LOGIN", response);
	if (!isPositiveIntermediate(status)) throw SMTPException("Cannot authenticate using LOGIN", response, status);

	std::ostringstream usernameBase64;
	CBase64Encode usernameEncoder(usernameBase64);
	usernameEncoder << username;
	usernameEncoder.close();

	std::ostringstream passwordBase64;
	CBase64Encode passwordEncoder(passwordBase64);
	passwordEncoder << password;
	passwordEncoder.close();

	//Server request for username/password not defined could be either
	//S: login:
	//C: user_login
	//S: password:
	//C: user_password
	//or
	//S: password:
	//C: user_password
	//S: login:
	//C: user_login

	std::string decodedResponse;
	std::istringstream responseStream(response.substr(4));
	CBase64Decode responseDecoder(responseStream);
	CStreamCopier::copyToString(responseDecoder, decodedResponse);

	if (icompare(decodedResponse, 0, 8, "username") == 0) // username first (md5("Username:"))
	{
		status = sendCommand(usernameBase64.str(), response);
		if (!isPositiveIntermediate(status)) throw SMTPException("Login using LOGIN username failed", response, status);

		status = sendCommand(passwordBase64.str(), response);
		if (!isPositiveCompletion(status)) throw SMTPException("Login using LOGIN password failed", response, status);  
	}
	else if  (icompare(decodedResponse, 0, 8, "password") == 0) // password first (md5("Password:"))
	{
		status = sendCommand(passwordBase64.str(), response);
		if (!isPositiveIntermediate(status)) throw SMTPException("Login using LOGIN password failed", response, status);  

		status = sendCommand(usernameBase64.str(), response);
		if (!isPositiveCompletion(status)) throw SMTPException("Login using LOGIN username failed", response, status);
	}
}

void CSMTPClientSession::loginUsingPlain(const std::string& username, const std::string& password)
{
	std::ostringstream credentialsBase64;
	CBase64Encode credentialsEncoder(credentialsBase64);
	credentialsEncoder << '\0' << username << '\0' << password;
	credentialsEncoder.close();

	std::string response;
	int status = sendCommand("AUTH PLAIN", credentialsBase64.str(), response);
	if (!isPositiveCompletion(status)) throw SMTPException("Login using PLAIN failed", response, status);
}

void CSMTPClientSession::login(LoginMethod loginMethod, const std::string& username, const std::string& password)
{
	login(CEnvironment::nodeName(), loginMethod, username, password);
}

void CSMTPClientSession::login(const std::string& hostname, LoginMethod loginMethod, const std::string& username, const std::string& password)
{
	std::string response;
	login(hostname, response);

	if (loginMethod == AUTH_CRAM_MD5)
	{
		if (response.find("CRAM-MD5", 0) != std::string::npos)
		{
			loginUsingCRAMMD5(username, password);
		}
		else throw SMTPException("The mail service does not support CRAM-MD5 authentication", response);
	}
	else if (loginMethod == AUTH_CRAM_SHA1)
	{
		if (response.find("CRAM-SHA1", 0) != std::string::npos)
		{
			loginUsingCRAMSHA1(username, password);
		}
		else throw SMTPException("The mail service does not support CRAM-SHA1 authentication", response);
	}
	else if (loginMethod == AUTH_LOGIN)
	{
		if (response.find("LOGIN", 0) != std::string::npos)
		{
			loginUsingLogin(username, password);
		}
		else throw SMTPException("The mail service does not support LOGIN authentication", response);
	}
	else if (loginMethod == AUTH_PLAIN)
	{
		if (response.find("PLAIN", 0) != std::string::npos)
		{
			loginUsingPlain(username, password);
		}
		else throw SMTPException("The mail service does not support PLAIN authentication", response);
	}
	else if (loginMethod != AUTH_NONE)
	{
		throw SMTPException("The autentication method is not supported");
	}
}

void CSMTPClientSession::open()
{
	if (!m_isOpen)
	{
		std::string response;
		int status = m_socket.receiveStatusMessage(response);
		if (!isPositiveCompletion(status)) throw SMTPException("The mail service is unavailable", response, status);
		m_isOpen = true;
	}
}

void CSMTPClientSession::close()
{
	if (m_isOpen)
	{
		std::string response;
		sendCommand("QUIT", response);
		m_socket.close();
		m_isOpen = false;
	}
}

void CSMTPClientSession::sendCommands(const CMailMessage& message, const Recipients* pRecipients)
{
	std::string response;
	int status = 0;
	const std::string& fromField = message.getSender();
	std::string::size_type emailPos = fromField.find('<');
	if (emailPos == std::string::npos)
	{
		std::string sender("<");
		sender.append(fromField);
		sender.append(">");
		status = sendCommand("MAIL FROM:", sender, response);
	}
	else
	{
		status = sendCommand("MAIL FROM:", fromField.substr(emailPos, fromField.size() - emailPos), response);
	}

	if (!isPositiveCompletion(status)) throw SMTPException("Cannot send message", response, status);

	std::ostringstream recipient;
	if (pRecipients)
	{
		for (Recipients::const_iterator it = pRecipients->begin(); it != pRecipients->end(); ++it)
		{
			recipient << '<' << *it << '>';
			int status = sendCommand("RCPT TO:", recipient.str(), response);
			if (!isPositiveCompletion(status)) throw SMTPException(std::string("Recipient rejected: ") + recipient.str(), response, status);
			recipient.str("");
		}
	}
	else
	{
		for (CMailMessage::Recipients::const_iterator it = message.recipients().begin(); it != message.recipients().end(); ++it)
		{
			recipient << '<' << it->getAddress() << '>';
			int status = sendCommand("RCPT TO:", recipient.str(), response);
			if (!isPositiveCompletion(status)) throw SMTPException(std::string("Recipient rejected: ") + recipient.str(), response, status);
			recipient.str("");
		}
	}

	status = sendCommand("DATA", response);
	if (!isPositiveIntermediate(status)) throw SMTPException("Cannot send message data", response, status);
}

void CSMTPClientSession::sendAddresses(const std::string& from, const Recipients& recipients)
{
	std::string response;
	int status = 0;

	std::string::size_type emailPos = from.find('<');
	if (emailPos == std::string::npos)
	{
		std::string sender("<");
		sender.append(from);
		sender.append(">");
		status = sendCommand("MAIL FROM:", sender, response);
	}
	else
	{
		status = sendCommand("MAIL FROM:", from.substr(emailPos, from.size() - emailPos), response);
	}

	if (!isPositiveCompletion(status)) throw SMTPException("Cannot send message", response, status);

	std::ostringstream recipient;

	for (Recipients::const_iterator it = recipients.begin(); it != recipients.end(); ++it)
	{

		recipient << '<' << *it << '>';
		int status = sendCommand("RCPT TO:", recipient.str(), response);
		if (!isPositiveCompletion(status)) throw SMTPException(std::string("Recipient rejected: ") + recipient.str(), response, status);
		recipient.str("");
	}
}

void CSMTPClientSession::sendData()
{
	std::string response;
	int status = sendCommand("DATA", response);
	if (!isPositiveIntermediate(status)) throw SMTPException("Cannot send message data", response, status);
}

void CSMTPClientSession::sendMessage(const CMailMessage& message)
{
	sendCommands(message);
	transportMessage(message);
}

void CSMTPClientSession::sendMessage(const CMailMessage& message, const Recipients& recipients)
{
	sendCommands(message, &recipients);
	transportMessage(message);
}

void CSMTPClientSession::transportMessage(const CMailMessage& message)
{
	CSocketOutputStream socketStream(m_socket);
	CMailOutputStream mailStream(socketStream);
	message.write(mailStream);
	mailStream.close();
	socketStream.flush();

	std::string response;
	int status = m_socket.receiveStatusMessage(response);
	if (!isPositiveCompletion(status)) throw SMTPException("The server rejected the message", response, status);
}

int CSMTPClientSession::sendCommand(const std::string& command, std::string& response)
{
	m_socket.sendMessage(command);
	return m_socket.receiveStatusMessage(response);
}

int CSMTPClientSession::sendCommand(const std::string& command, const std::string& arg, std::string& response)
{
	m_socket.sendMessage(command, arg);
	return m_socket.receiveStatusMessage(response);
}

void CSMTPClientSession::sendMessage(std::istream& istr)
{
	std::string response;
	int status = 0;

	CSocketOutputStream socketStream(m_socket);
	CMailOutputStream mailStream(socketStream);
	CStreamCopier::copyStream(istr, mailStream);
	mailStream.close();
	socketStream.flush();
	status = m_socket.receiveStatusMessage(response);
	if (!isPositiveCompletion(status)) throw SMTPException("The server rejected the message", response, status);
}

bool CSMTPClientSession::isPositiveCompletion(int status)
{
	return status/100 == SMTP_POSITIVE_COMPLETION;
}

bool CSMTPClientSession::isPositiveIntermediate(int status)
{
	return status/100 == SMTP_POSITIVE_INTERMEDIATE;
}

bool CSMTPClientSession::isTransientNegative(int status)
{
	return status/100 == SMTP_TRANSIENT_NEGATIVE;
}

bool CSMTPClientSession::isPermanentNegative(int status)
{
	return status/100 == SMTP_PERMANENT_NEGATIVE;
}

CDialogSocket& CSMTPClientSession::socket()
{
	return m_socket;
}

NET_NAMESPACE_END
